import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Activation
from tensorflow.keras.layers import Dense
from tensorflow.keras import utils
# 时间步：一个字母一个时间步

sample = "hihello"

# 模型只能计算数字，字母转化为数字
char_set = list(set(sample))  # id -> char ['i', 'l', 'e', 'o', 'h']
char_dic = {w: i for i, w in enumerate(char_set)}

# x: hihell
# y: ihello
x_str = sample[:-1]
y_str = sample[1:]

data_dim = len(char_set) # 有多少个字母
timesteps = len(y_str) # 时间步， y标签长度， 每个字母一个时间步
num_classes = len(char_set) # 预测每个字符的结果

print(x_str, y_str)

x = [char_dic[c] for c in x_str]  # 字符转化为数字形式  [6]
y = [char_dic[c] for c in y_str]  # char to index

# One-hot encoding
x = utils.to_categorical(x, num_classes=num_classes) # [6, 5]
# reshape X to be [samples, time steps, features]
x = np.reshape(x, (-1, len(x), data_dim))
print(x.shape) # [ 1, 6, 5]

# One-hot encoding
y = utils.to_categorical(y, num_classes=num_classes)
# time steps
y = np.reshape(y, (-1, len(y), data_dim))
print(y.shape)

model = Sequential()
model.add(Dense(num_classes, input_shape=(
    timesteps, data_dim)))
model.add(Activation('softmax')) # 全连接的多分类
model.summary()

model.compile(loss='categorical_crossentropy',
              optimizer='rmsprop', metrics=['accuracy'])
model.fit(x, y, epochs=10000)

predictions = model.predict(x, verbose=0)# [1, 6, 5]
print('---\n', predictions)
for i, prediction in enumerate(predictions): # prediction [6， 5]
    x_index = np.argmax(x[i], axis=1)
    x_str = [char_set[j] for j in x_index]
    print(x_index, ''.join(x_str))

    index = np.argmax(prediction, axis=1)
    result = [char_set[j] for j in index]
    print(index, ''.join(result))
